import torch
import pandas as pd
import pickle
from utils import TextDataset
from torch.utils.data import DataLoader
from models.attention_rnn_model import AttentionRNN  # 或换成你使用的模型

# 超参数必须和训练时一致
embed_dim = 100
hidden_dim = 128
num_classes = 2
batch_size = 16

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载 vocab
with open('dataset/vocab.pkl', 'rb') as f:
    vocab = pickle.load(f)

# 加载模型
model = AttentionRNN(len(vocab), embed_dim, hidden_dim, num_classes).to(device)
model.load_state_dict(torch.load('model.pth', map_location=device))
model.eval()

# 加载测试数据
test_data = pd.read_csv('dataset/test.csv')
test_dataset = TextDataset(list(zip(test_data['text'], test_data['label'])), vocab)
test_loader = DataLoader(test_dataset, batch_size=batch_size)

# 评估准确率
correct = total = 0
with torch.no_grad():
    for x, y in test_loader:
        x, y = x.to(device), y.to(device)
        output = model(x)
        _, predicted = output.max(1)
        total += y.size(0)
        correct += predicted.eq(y).sum().item()

print(f"✅ 测试集准确率: {correct / total:.4f}")
